{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-0BQWhvAP2jb"
      },
      "source": [
        "\n",
        "\u003ca href=\"https://colab.research.google.com/github/google-research/t5x/blob/main/t5x/notebooks/introduction.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bqZYp90PIa1t"
      },
      "source": [
        "# Overview\n",
        "\n",
        "T5X is a modular, composable, research-friendly framework for high-performance, configurable, self-service training, evaluation, and inference of sequence models (starting with language) at many scales.\n",
        "\n",
        "It is essentially a new and improved implementation of the [T5 codebase](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md) (based on Mesh TensorFlow) in JAX and Flax.\n",
        "\n",
        "# Getting Started\n",
        "\n",
        "In the following Colab, we present an introductory tutorial to get you started interacting with the T5X codebase. In particular, we'll introduce the major components of the T5X codebase and get you started running training, inference, and evaluation on natural text inputs.\n",
        "\n",
        "Note: If you are a using public colab, please use its `Connect to a local runtime` option by following the [setup guide](https://github.com/google-research/t5x/blob/main/t5x/notebooks/README.md)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jIGSIHzD7YPO"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "import os\n",
        "\n",
        "import clu.data.dataset_iterator\n",
        "import tensorflow as tf\n",
        "import jax\n",
        "from jax import random\n",
        "from jax.experimental import multihost_utils\n",
        "import jax.numpy as jnp\n",
        "from flax import linen\n",
        "import numpy as np\n",
        "import seqio\n",
        "import t5.data\n",
        "from t5.evaluation import metrics as t5_metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q2NtKT-vlomE"
      },
      "outputs": [],
      "source": [
        "import t5x\n",
        "from t5x import partitioning\n",
        "from t5x import train_state as train_state_lib\n",
        "from t5x import utils\n",
        "from t5x.examples.t5 import network\n",
        "from t5x.examples.scalable_t5 import network as scalable_network\n",
        "from t5x.interactive_model import InteractiveModel\n",
        "from t5x.interactive_model import get_batches_from_seqio\n",
        "from t5x.interactive_model import InferenceType"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ib9aOi2xaCKQ"
      },
      "source": [
        "# T5X Components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-trKAuOWaGoB"
      },
      "source": [
        "Let's start by going over some of the major components of the T5X codebase: models, checkpoints, and partitioners.\n",
        "\n",
        "We will define instances of some of these components in the following subsections before we use them to run training, inference, and evaluation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RE_CQr9Hcr1D"
      },
      "source": [
        "## T5X Models\n",
        "One of the primary contributions of the T5X codebase is its easy-to-use collection of models.\n",
        "\n",
        "The T5X codebase provides an abstract base class, [`BaseModel`](https://github.com/google-research/t5x/blob/main/t5x/models.py), which should be subclassed to define specific model architectures. This abstraction allows us to flexibly extend the T5X framework to custom architectures. Importantly, the `BaseModel` and all subclasses are free from parallelism-related features (this is handled by the partitioner; see following sections).\n",
        "\n",
        "The T5X codebase also provides several widely-used subclasses of the `BaseModel`, namely the [`EncoderDecoderModel`](https://github.com/google-research/t5x/blob/main/t5x/models.py) and the [`DecoderModel`](https://github.com/google-research/t5x/blob/main/t5x/models.py).\n",
        "\n",
        "Importantly, the proposed structure of the `BaseModel`/all subclasses does not impose that the model be implemented in a specific framework. Instead, all subclasses of the `BaseModel` take in an `nn.Module` constructor argument, which is used to implement the architecture of the model. These modules can be built in Flax (e.g. [minimal T5](https://github.com/google-research/t5x/blob/main/t5x/examples/t5/network.py)) or on top of a layers library such as [Flaxformer](https://github.com/google/flaxformer).\n",
        "\n",
        "We've provided a sample model definition below. For this example, we will instantiate an `EncoderDecoderModel`, which will also require us to define input and output vocabularies, an optimizer, and a decode function. We'll use the minimal T5 module to implement our model architecture."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FqH_jEcUdEsd"
      },
      "outputs": [],
      "source": [
        "# Define EncoderDecoderModel constructor args (except the module).\n",
        "input_vocabulary=t5.data.get_default_vocabulary()\n",
        "output_vocabulary=t5.data.get_default_vocabulary()\n",
        "optimizer=t5x.adafactor.Adafactor(decay_rate=0.8, step_offset=0, logical_factor_rules=t5x.adafactor.standard_logical_factor_rules())\n",
        "decode_fn=functools.partial(t5x.decoding.temperature_sample, temperature=1.0, topk=40)\n",
        "\n",
        "# Define a model using the minimal T5 module.\n",
        "t5_module = network.Transformer(config=network.T5Config(\n",
        "    vocab_size=32128,\n",
        "    dtype='bfloat16',\n",
        "    emb_dim=512,\n",
        "    num_heads=6,\n",
        "    num_encoder_layers=8,\n",
        "    num_decoder_layers=8,\n",
        "    head_dim=64,\n",
        "    mlp_dim=1024,\n",
        "    mlp_activations=('gelu', 'linear'),\n",
        "    dropout_rate=0.0,\n",
        "    logits_via_embedding=False))\n",
        "model = t5x.models.EncoderDecoderModel(\n",
        "    module=t5_module,\n",
        "    input_vocabulary=input_vocabulary,\n",
        "    output_vocabulary=output_vocabulary,\n",
        "    optimizer_def=optimizer,\n",
        "    decode_fn=decode_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r74x8nJpfe3G"
      },
      "source": [
        "## Checkpoints\n",
        "\n",
        "The T5X codebase also includes checkpoints for a wide variety of pre-trained T5X models. A full list of all publically available checkpoints is available at https://github.com/google-research/t5x/blob/main/docs/models.md.\n",
        "\n",
        "For the following example, we have selected a pretrained [T5 1.1 Small model](https://github.com/google-research/t5x/blob/main/docs/models.md) that has been additionally finetuned to answer natural questions using the (open domain) [Natural Questions benchmark](https://ai.google.com/research/NaturalQuestions). We use this finetuned checkpoint for this example in order to see improved performance on the natural question examples we will use for training/inference/evaluation later on.\n",
        "\n",
        "To restore our model from this checkpoint, we first define the path to our checkpoint and the `dtype` to restore."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvVNR5AWgPOC"
      },
      "outputs": [],
      "source": [
        "# The checkpoint below is a T5-1.1-Small checkpoint (https://github.com/google-research/t5x/blob/main/docs/models.md)\n",
        "# that has additionally been finetuned on the (Open Domain) Natural Questions\n",
        "# benchmark (https://ai.google.com/research/NaturalQuestions).\n",
        "checkpoint_path='gs://t5-data/pretrained_models/cbqa/small_ssm_nq/model.ckpt-1110000'\n",
        "dtype='bfloat16'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Obgf5nfMgdkm"
      },
      "source": [
        "We also need to define how we want to restore our model. There are two different restore modes that are available in T5X; for now, we will use \"specific\", which will load the most recent checkpoint in the directory specified by `checkpoint_path`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eqfWb6GZhVz-"
      },
      "outputs": [],
      "source": [
        "restore_mode='specific'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PkG5uRfdhnzR"
      },
      "source": [
        "Finally, it should be noted that if you are restoring your model from a checkpoint, then the model architecture you defined above must match the model architecture of your checkpoint. For all T5X checkpoints listed at https://github.com/google-research/t5x/blob/main/docs/models.md, you can find the correct architecture for the given checkpoint in its corresponding Gin file."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4wzrzzTiRbl"
      },
      "source": [
        "## Partitioners\n",
        "\n",
        "Partitioning is the process of dividing and replicating machine learning model parameters, activations, and data across accelerator devices in order to:\n",
        "\n",
        "\n",
        "*   Train and infer from models too large to fit in the memory of a single device\n",
        "*   Use extremely large batch sizes\n",
        "*   Train faster\n",
        "\n",
        "In T5X, partitioning is primarily provided through the [jax.pjit](https://github.com/google/jax/tree/main/jax/experimental/pjit.py) fronted via `PjitPartitioner`. `PjitPartitioner` has three primary constructor arguments:\n",
        "*    `model_parallel_submesh`\n",
        "*    `num_partitions`\n",
        "*    `logical_axis_rules`\n",
        "\n",
        "The `model_parallel_submesh` and `num_partitions` arugments provide two mutually exclusive methods of specifying the submesh of devices to use for model partitioning. If you specify `num_partitions`, T5X will use this value to generate a default `model_parallel_submesh` that is suitable, but may not be the optimal configuration. If you are interested in optimizing performance, you can try out different submeshes using the `model_parallel_submesh` parameter. For simplicity, we will use `num_partitions` in this Colab.\n",
        "\n",
        "If you are interested in learning more about partitioning, please take a look at our T5X: Partitioning Deep Dive Colab (Colab status: WIP, link is upcoming).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kWEmFrvRkAs6"
      },
      "outputs": [],
      "source": [
        "partitioner=partitioning.PjitPartitioner(\n",
        "        num_partitions=1,\n",
        "        model_parallel_submesh=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8b1ECSqVlCcl"
      },
      "source": [
        "# Running Training, Inference, and Evaluation\n",
        "\n",
        "Now, let's get started running training, inference, and evaluation on natural text inputs. T5X provides an `InteractiveModel` class that we can wrap around our model, checkpoint, and partitioner components, enabling us to run training, inference, and evaluation in one line of code each.\n",
        "\n",
        "The InteractiveModel requires a couple of additional constructor arguments, namely:\n",
        "\n",
        "\n",
        "1.   `batch_size`: the number of examples per batch for training, inference, and evaluation.\n",
        "2.   `task_feature_lengths`:  `task_feature_lengths` is a dictionary mapping the task feature key to the maximum length (int) for that feature. If a feature is longer than this length after preprocessing, the feature will be truncated. May be set to `None` to avoid truncation. \\\n",
        "For context, task features are specific to tasks (ex: inputs and targets), and can be mapped to various model-specific features (for example, if we are using a decoder-only model, the concatenation of inputs and targets will be mapped to `decoder_target_tokens`, the model features). This mapping is done by the model's feature converter.\n",
        "3.   `output_dir`: Path to directory where we will write new model checkpoints.\n",
        "4.   `input_shapes`: a mapping from key to array shape for each model feature in the global (unsharded) input batch. These input shapes are used to define and initialize the train state. Importantly, these input shapes define the *model features* shape, in contrast to the task features described above.\n",
        "\n",
        "We define these arguments and an instance of the InteractiveModel below. Importantly, it should be noted that the InteractiveModel handles restoring our model from the provided checkpoint path, so once we instantiate the InteractiveModel, we will be ready to run training, inference, and evaluation. Restoring the model from a checkpoint may take a minute or two.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ARvXFgR6l8T5"
      },
      "outputs": [],
      "source": [
        "batch_size=8\n",
        "task_feature_lengths = {'inputs': 38, 'targets': 18}\n",
        "output_dir='/tmp/output_dir'\n",
        "input_shapes = {\n",
        "    'encoder_input_tokens': np.array([8, 38]),\n",
        "    'decoder_target_tokens': np.array([8, 18]),\n",
        "    'decoder_input_tokens': np.array([8, 18]),\n",
        "    'decoder_loss_weights': np.array([8, 18])\n",
        "}\n",
        "\n",
        "interactive_model = InteractiveModel(\n",
        "  batch_size=batch_size,\n",
        "  task_feature_lengths=task_feature_lengths,\n",
        "  output_dir=output_dir,\n",
        "  partitioner=partitioner,\n",
        "  model=model,\n",
        "  dtype=dtype,\n",
        "  restore_mode=restore_mode,\n",
        "  checkpoint_path=checkpoint_path,\n",
        "  input_shapes=input_shapes\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOCiFtbyl2wc"
      },
      "source": [
        "Next, let's define some examples that we want to use for training/inference/evaluation. These examples should either be a list of inputs, or a list of dictionaries mapping 'target'/'input' keys to corresponding values, as shown below. We will define two sets of examples: one set to be trained on, and one set to run inference/evaluation on.\n",
        "\n",
        "We are using natural question/answer pairs for our examples. As described in the [T5 paper](https://arxiv.org/abs/1910.10683), we must add a task-specific prefix to our input before we feed it to the model in order to specify what task we should perform on the provided input. For natural questions, we use the \"nq question:\" prefix."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KouYfpKflw04"
      },
      "outputs": [],
      "source": [
        "training_examples = [\n",
        "  {\n",
        "      'input': 'nq question: who has been appointed as the new chairman of sebi',\n",
        "      'target': 'Ajay Tyagi'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: who wrote the book lion the witch and the wardrobe',\n",
        "      'target': 'C. S. Lewis'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: how many planes did japan lose at pearl harbor',\n",
        "      'target': '29'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: who does the voice of mcgruff the dog',\n",
        "      'target': 'Jack Keil'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: who sings the wheels in the sky keep on turning',\n",
        "      'target': 'Journey'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: who voices regina in glitter force doki doki',\n",
        "      'target': 'Kumiko Watanabe'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: when did the us become allies with britain',\n",
        "      'target': 'during World War II'\n",
        "  },\n",
        "  {\n",
        "      'input': 'nq question: who won the rugby 7 in las vegas',\n",
        "      'target': 'the United States'\n",
        "  },\n",
        "]\n",
        "\n",
        "validation_examples = [\n",
        "  {\n",
        "      'target': 'Joe Biden',\n",
        "      'input':'nq question: who is the president of the united states'\n",
        "  },\n",
        "  {\n",
        "      'target': 'F. Scott Fitzgerald',\n",
        "      'input': 'nq question: who wrote the book the great gatsby'},\n",
        "  {\n",
        "      'target': '1914',\n",
        "      'input': 'nq question: in what year did the first world war begin'},\n",
        "  {\n",
        "      'target': 'Idina Menzel',\n",
        "      'input': 'nq question: who does the voice of elsa in Frozen'},\n",
        "  {\n",
        "      'target': 'Taylor Swift',\n",
        "      'input': 'nq question: who sings shake it off'},\n",
        "  {\n",
        "      'target': 'Tom Kenny',\n",
        "      'input': 'nq question: who voices spongebob squarepants'},\n",
        "  {\n",
        "      'target': '2010',\n",
        "      'input': 'nq question: when did the great british bake off start'},\n",
        "  {\n",
        "      'target': 'the Philadelphia Eagles',\n",
        "      'input': 'nq question: who won the superbowl in 2018'},\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAMJOydml0NU"
      },
      "source": [
        "Now, we can run training, inference and evaluation on these examples with a single line of code for each task. Below, we run training and inference (evaluation requires a few more arguments, so we go over evaluation in a following section). This may take ~60 seconds."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8OsCbMb0XlC0"
      },
      "outputs": [],
      "source": [
        "interactive_model.train_step(examples=training_examples)\n",
        "print(f\"Training Summary: {interactive_model.train_summary}\\n\")\n",
        "print(f\"Step Number: {interactive_model.step}\\n\")\n",
        "\n",
        "examples_and_predictions, _ = interactive_model.predict_with_aux(examples=validation_examples)\n",
        "predictions = [prediction for example, prediction in examples_and_predictions]\n",
        "print(f\"Predictions: {predictions}\\n\")\n",
        "\n",
        "examples_and_scores = interactive_model.score(examples=validation_examples)\n",
        "scores = [score for example, score in examples_and_scores]\n",
        "print(f\"Scores: {scores}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0SnCvC4d4pjA"
      },
      "source": [
        "Alternately, you can run a training/inference/evaluation loop over multiple batches. The training loop below runs training and inference for each step, using the provided batches, and returns the predictions and scores from the final step. This may take ~60 seconds (note: if you use XL or XXL model sizes, this loop may take a while to complete; we are working on improved compilation strategies that optimize for runtime in b/247170488)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GaBd-e5s4xZp"
      },
      "outputs": [],
      "source": [
        "second_batch_of_examples = [\n",
        "    {\n",
        "        'input': 'nq question: who won the most academy awards in his lifetime',\n",
        "        'target': 'Walt Disney'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: who starred in the hand that rocks the cradle',\n",
        "        'target': 'Rebecca De Mornay'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: what does a red license plate mean in ontario',\n",
        "        'target': 'diplomat'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: who sang i dreamed a dream on britain\\'s got talent',\n",
        "        'target': 'Susan Magdalane Boyle'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: when is season 7 of game of thrones being released',\n",
        "        'target': 'August 27, 2017'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: when is anne with an e season two coming out',\n",
        "        'target': 'in 2018'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: when was hard rock hotel las vegas built',\n",
        "        'target': '1995'\n",
        "    },\n",
        "    {\n",
        "        'input': 'nq question: what type of reaction leads to the production of polymers',\n",
        "        'target': 'condensation reaction'\n",
        "    }\n",
        "]\n",
        "all_training_batches = [training_examples, second_batch_of_examples]\n",
        "examples_and_predictions, examples_and_scores, _ = interactive_model.train_loop(num_steps=2, train_batches=all_training_batches, predict_batches=[validation_examples], score_batches=[validation_examples])\n",
        "\n",
        "print(\"\\n All Predictions\")\n",
        "for example, prediction in examples_and_predictions:\n",
        "  print(f\"Input: {example['inputs_pretokenized']}, Prediction: {prediction}\")\n",
        "print(\"\\nAll Scores:\")\n",
        "for example, score in examples_and_scores:\n",
        "  print(f\"Input: {example['inputs_pretokenized']}, Score: {score}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGaGfR84_8Ap"
      },
      "source": [
        "### Preprocessors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lpRZd2F4-G8l"
      },
      "source": [
        "By default, the only preprocessors that the methods above run are [tokenization](https://github.com/google/seqio/tree/main/seqio/preprocessors.py) and [appending an EOS token](https://github.com/google/seqio/tree/main/seqio/preprocessors.py). If you would like to use different preprocessors, you can do so using the `train_step_with_preprocessors` or `infer_with_preprocessors` methods. We've provided a sample below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dDWqZoeW_fRi"
      },
      "outputs": [],
      "source": [
        "preprocessors = [\n",
        "    seqio.preprocessors.tokenize,\n",
        "    seqio.preprocessors.append_eos\n",
        "]\n",
        "\n",
        "interactive_model.train_step_with_preprocessors(examples=training_examples, preprocessors=preprocessors)\n",
        "print(f\"Training Summary: {interactive_model.train_summary}\\n\")\n",
        "print(f\"Step Number: {interactive_model.step}\\n\")\n",
        "\n",
        "# Note: when we use a custom list of preprocessors, we must use a general\n",
        "# `infer` method, rather than `predict` or `score`. Thus, we must also specify\n",
        "# the type of inference to do; valid options are `PREDICT_WITH_AUX`,\n",
        "# or `SCORE`.\n",
        "examples_and_predictions, _ = interactive_model.infer_with_preprocessors(\n",
        "    mode=InferenceType.PREDICT_WITH_AUX,\n",
        "    examples=validation_examples,\n",
        "    preprocessors=preprocessors)\n",
        "predictions = [prediction for example, prediction in examples_and_predictions]\n",
        "print(f\"Predictions: {predictions}\\n\")\n",
        "\n",
        "examples_and_scores, _ = interactive_model.infer_with_preprocessors(\n",
        "    mode=InferenceType.SCORE,\n",
        "    examples=validation_examples,\n",
        "    preprocessors=preprocessors)\n",
        "scores = [score for example, score in examples_and_scores]\n",
        "print(f\"Scores: {scores}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sa08OJEuAuhv"
      },
      "source": [
        "Because we use the same set of preprocessors, we should expect to see the same results as before.\n",
        "\n",
        "If you are interested in learning more about preprocessors, please see [this preprocessors guide](https://github.com/google/seqio/blob/main/README.md#preprocessors), which also contains links to implementations of common preprocessors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNbOuE-kA21g"
      },
      "source": [
        "### Evaluation and Metrics Functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EXw52bCt8s3A"
      },
      "source": [
        "We can similarly run evaluation in a single line. Running evaluation requires that we specify a metric function and (optionally) a list of postprocessors to run on the data before we compute metrics.\n",
        "\n",
        "There are a variety of sample metrics defined in [t5/evaluation/metrics.py](https://github.com/google-research/text-to-text-transfer-transformer/tree/main/t5/evaluation/metrics.py). For this example, we will use the [SQuAD metric function](https://github.com/google-research/text-to-text-transfer-transformer/tree/main/t5/evaluation/metrics.py) defined in this file. Because we are using natural questions, we will also specify a postprocessor to correctly format question and answer pairs for metrics calculations; specifically, we will use a standard QA postprocessor, modeled after the [`t5.data.postprocessors.qa` method](https://github.com/google-research/text-to-text-transfer-transformer/tree/main/t5/data/postprocessors.py). We will continue to use the same preprocessors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_xr9iOEbxSzS"
      },
      "outputs": [],
      "source": [
        "def qa(answer, example=None, is_target=False):\n",
        "  \"\"\"Returns answer, or all answers if the full example is provided.\"\"\"\n",
        "  if is_target:\n",
        "    return [tf.compat.as_text(a) for a in [example[\"targets_pretokenized\"]]]\n",
        "  return answer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XLL9_3OwCU9i"
      },
      "outputs": [],
      "source": [
        "metrics = interactive_model.evaluate_with_preprocessors(\n",
        "        examples=validation_examples,\n",
        "        preprocessors=preprocessors,\n",
        "        metric_fns=[t5_metrics.squad],\n",
        "        postprocessor=qa)\n",
        "print(f\"Metrics: {metrics}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OJ5NHdcoPFVX"
      },
      "source": [
        "If you are interested in learning more about metrics functions or postprocessors, please see this [Metrics guide](https://github.com/google/seqio/blob/main/README.md#metrics ) and/or this [Postprocessors guide](https://github.com/google/seqio/blob/main/README.md#postprocessor)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lcDwmp_AxnOG"
      },
      "source": [
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-QR5LnmN4ikp"
      },
      "source": [
        "# Advanced Topics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wi29fMdv4mSr"
      },
      "source": [
        "## SeqIO\n",
        "\n",
        "If you are interested in T5X, you may also be interested in, or have heard of, SeqIO. SeqIO is a library for processing sequential data to be fed into downstream sequence models. At a high level, SeqIO relies on user-defined `Tasks` and `Mixtures` that can be used to retrieve and evaluate datasets.\n",
        "\n",
        "We won't go into details about SeqIO here; we recommend checking out this [SeqIO Introductory guide](https://github.com/google/seqio/blob/main/README.md) and/or clicking below to run a SeqIO Introductory Colab. The rest of this section will assume a basic understanding of SeqIO.\n",
        "\n",
        "\u003ca href=\"https://colab.research.google.com/github/google/seqio/blob/main/seqio/notebooks/Basics_Task_and_Mixtures.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "If you are already familiar with SeqIO and have a SeqIO task/mixture that you would like to use in this Colab, we do provide a SeqIO bridge that takes in a SeqIO task/mixture and produces batches of examples that can be processed by the InteractiveModel above. We've provided an example of this bridge below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hxu9mRL5yBGK"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/google-research/google-research.git google_research"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bM0nRIEFwyj_"
      },
      "outputs": [],
      "source": [
        "import google_research.t5_closed_book_qa.t5_cbqa.tasks\n",
        "batches = get_batches_from_seqio(\n",
        "        task_or_mixture_name='natural_questions_open',\n",
        "        split='validation',\n",
        "        batch_size=8,\n",
        "        num_batches=2,\n",
        "        seed=42)\n",
        "print(f\"Batches: {batches}\")\n",
        "# Train the interactive model on the provided batches.\n",
        "original_step = interactive_model.step\n",
        "_ = interactive_model.train_loop(num_steps=len(batches), train_batches=batches)\n",
        "print(f\"Original Step: {original_step}, Current Step: {interactive_model.step}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Elt08160w03X"
      },
      "source": [
        "The `get_batches_from_seqio` bridge can take several constructor arguments:\n",
        "\n",
        "\n",
        "1.   `task_or_mixture_name`: the name of the SeqIO task/mixture to read data from. It should be noted that your task/mixture must already be registered with SeqIO, and you must import the module that defines your task/mixture here (as seen above).\n",
        "2.   `split`: the split of the Task/Mixture to read data from.\n",
        "3.   `batch_size`: how many examples should appear in each batch.\n",
        "4.   `num_batches`: the total number of batches to return.\n",
        "5.   `get_pretokenized_examples`: optional. A boolean, defaulting to True, that determines whether we should read the `inputs_pretokenized`/`targets_pretokenized` elements from an example, or the `inputs`/`targets` elements. \\\n",
        "The `train_step`, `predict`, `predict_with_aux`, `score`, and `evaluate` methods of the InteractiveModel assume that we should run [tokenization](https://github.com/google/seqio/tree/main/seqio/preprocessors.py) and [appending an EOS token](https://github.com/google/seqio/tree/main/seqio/preprocessors.py) as the only preprocessors. To use these methods with this pre-defined list of preprocessors, you can set `get_pretokenized_examples=True` to retrieve examples that still need to be tokenized, and these InteractiveModel methods will handle running these preprocessors. This setting can also be helpful if you want to inspect the natural text inputs/targets of your SeqIO task. \\\n",
        "However, some SeqIO tasks do not use tokenization (ex: span corruption). You can set `get_pretokenized_examples=False`, and this bridge will read the fully preprocessed examples from the SeqIO task. You can then run `train_step_with_preprocessors`, `infer_with_preprocessors`, or `evaluate_with_preprocessors` and provide an empty preprocessors list (because all preprocessing has already been completed by this bridge) to run training/inference/evaluation. We have provided an example of using this bridge to retrieve fully preprocessed examples below.\n",
        "6.   `sequence_length`: optional. A dictionary mapping feature key to maximum length (int) for that feature. Used by SeqIO to retrieve the dataset/examples.\n",
        "7.   `**get_dataset_kwargs`: there are many [additional parameters](https://github.com/google/seqio/tree/main/seqio/dataset_providers.py) that can be set in the `SeqIO.get_dataset` function. If you would like to set any of these arguments, you can set them using this `kwargs` parameter.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fjKBCX39w0Xl"
      },
      "outputs": [],
      "source": [
        "import t5.data.tasks\n",
        "batches = get_batches_from_seqio(\n",
        "    task_or_mixture_name='c4_v220_span_corruption',\n",
        "    split='validation',\n",
        "    batch_size=8,\n",
        "    num_batches=1,\n",
        "    get_pretokenized_examples=False,\n",
        "    sequence_length=interactive_model._task_feature_lengths,\n",
        "    seed=42)\n",
        "batch = batches[0]  # We expect only a single batch.\n",
        "original_step = interactive_model.step\n",
        "interactive_model.train_step_with_preprocessors(\n",
        "        examples=batch, preprocessors=[])\n",
        "print(f\"Original Step: {original_step}, Current Step: {interactive_model.step}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "RE_CQr9Hcr1D",
        "r74x8nJpfe3G",
        "a4wzrzzTiRbl"
      ],
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "Welcome to T5X: An Introductory Colab",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Akpc6pKlJB5rn5YYYFC9lw2OMk6oBzlQ",
          "timestamp": 1662951101563
        },
        {
          "file_id": "1rA8bgO2bJRoebAuS96Ji0RUhnawgBY4i",
          "timestamp": 1650477076639
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
